from app.db import db
from app.models.user_model import User
from app.models.article_model import Article

def init_db():
    """初始化数据库连接"""
    db.connect()
    # 在这里可以添加创建表的逻辑
    print("数据库连接成功")
    # 检查表是否存在，如果不存在则创建
    if not User.table_exists():
        with db:
            db.create_tables([User])
    # 检查文章表是否存在，如果不存在则创建
    if not Article.table_exists():
        with db:
            db.create_tables([Article])        
    db.close()

def get_db():
    """获取数据库连接"""
    return db